# gc();rm(list = ls());
gc()

# library(dplyr)
# library(ggplot2)
# library(tidyr)
# library(forcats)

dir.create("plots/us/", recursive = TRUE)

load("results/us/full.Rda")

supp.labs <- c("Inter-party Dimension", "Intra-party Dimension")
names(supp.labs) <- c("1", "2")
coef.labs <- c("Group Indicator", "Ideal Point (t-1)")
names(coef.labs) <- c("2", "3")

# nonseparability plot
ggplot(ns_est_full,
       aes(x = session, y = mean, ymin = q5, ymax = q95)) + geom_point() + geom_linerange() + theme_bw() +
  xlab("Congress") + ylab("Non-separability") + geom_hline(aes(yintercept = 0), linetype = 2, size = .5) +
  geom_smooth(aes(linetype = parameter), color = "black", size = .5, alpha = .2) +
  theme(legend.position = "bottom", legend.title = element_blank()) + scale_linetype_discrete(labels = c("Non-separable")) + scale_shape_discrete(labels = c("Non-separable"))
ggsave(filename = "plots/us/nsep_e.pdf", width = 4, height = 4)

# salience plot 
sal_est_full %>% select(mean, q5, q95, model, id, session) %>% 
  ggplot(aes(x = session, y = mean, ymin = q5, ymax = q95, shape = model)) + geom_point(size = 2, position = position_dodge(width = .3)) + geom_linerange(position = position_dodge(width = .3)) + theme_bw() +
  xlab("Congress") + ylab("Salience") + facet_wrap(~id, ncol = 1, labeller = labeller(id = supp.labs)) + 
  geom_smooth(aes(linetype = model), position = position_dodge(width = .3), color = "black", size = .5, alpha = .2) + 
  theme(legend.position = "bottom", legend.title = element_blank()) + scale_linetype_discrete(labels = c("Non-separable", "Separable")) + scale_shape_discrete(labels = c("Non-separable", "Separable"))
ggsave(filename = "plots/us/sal_e.pdf", width = 4, height = 6)

weights_est_full %>% select(mean, q5, q95, model, id, session) %>% filter(id %in% c(1, 4)) %>% 
  ggplot(aes(x = session, y = mean, ymin = q5, ymax = q95, shape = model)) + geom_point(size = 2, position = position_dodge(width = .3)) + geom_linerange(position = position_dodge(width = .3)) + theme_bw() +
  xlab("Congress") + ylab("Salience") + facet_wrap(~id, ncol = 1, labeller = labeller(id = supp.labs)) + 
  geom_smooth(aes(linetype = model), position = position_dodge(width = .3), color = "black", size = .5, alpha = .2) + 
  theme(legend.position = "bottom", legend.title = element_blank()) + scale_linetype_discrete(labels = c("Non-separable", "Separable")) + scale_shape_discrete(labels = c("Non-separable", "Separable"))

# correlation plot
cor_est_full %>% select(mean, q5, q95, model, session) %>% 
  ggplot(aes(x = session, y = mean, ymin = q5, ymax = q95, shape = model)) + geom_point(size = 2, position = position_dodge(width = .3)) + geom_linerange(position = position_dodge(width = .3)) + theme_bw() +
  xlab("Congress") + ylab("Correlation across Dimensions") + geom_smooth(aes(linetype = model), position = position_dodge(width = .3), color = "black", size = .5, alpha = .2) + 
  geom_hline(aes(yintercept = 0), linetype = 2) + 
  theme(legend.position = "bottom", legend.title = element_blank()) + scale_linetype_discrete(labels = c("Non-separable", "Separable")) + scale_shape_discrete(labels = c("Non-separable", "Separable"))
ggsave(filename = "plots/us/cor_e.pdf", width = 4, height = 4)

# correlation plot
old_est_full %>% select(mean, q5, q95, model, id, session) %>% filter(session > 79) %>% 
  ggplot(aes(x = session, y = mean, ymin = q5, ymax = q95, shape = model)) + geom_point(position = position_dodge(width = .3)) + geom_linerange(position = position_dodge(width = .3)) + theme_bw() +
  xlab("Congress") + ylab("Correlation over Time") + facet_wrap(~id, ncol = 1, labeller = labeller(id = supp.labs)) + 
  geom_smooth(aes(linetype = model), position = position_dodge(width = .3), color = "black", size = .5, alpha = .2) + 
  theme(legend.position = "bottom", legend.title = element_blank()) + scale_linetype_discrete(labels = c("Non-separable", "Separable")) + scale_shape_discrete(labels = c("Non-separable", "Separable"))
ggsave(filename = "plots/us/corold_e.pdf", width = 6, height = 6)

old_est_full %>% select(mean, model, id, session) %>% filter(session > 79) %>% 
  group_by(id) %>% 
  summarize(m = mean(mean),
            sd = sd(mean))

old_est_full %>% select(mean, model, id, session) %>% filter(session > 79 & id == 2) %>% 
  arrange(mean)

# theta plot
id_est_full %>% select(mean, model, id, dim, session) %>% 
  pivot_wider(names_from = dim, values_from = mean) %>% 
  rename(d1 = `1`, d2 = `2`) %>% 
  filter(model == "sep") %>% 
  ggplot(aes(d1, d2))+geom_point(alpha = .1) + xlab("Inter-party Dimension") + ylab("Intra-party Dimension") + theme_bw() + facet_wrap(~session, ncol = 5) + 
  geom_smooth(method = "lm", color = "black", alpha = .2, size = .5)
ggsave(filename = "plots/us/ideal_sep_e.pdf", width = 6, height = 8)
id_est_full %>% select(mean, model, id, dim, session) %>% 
  pivot_wider(names_from = dim, values_from = mean) %>% 
  rename(d1 = `1`, d2 = `2`) %>% 
  filter(model != "sep") %>% 
  ggplot(aes(d1, d2))+geom_point(alpha = .1) + xlab("Inter-party Dimension") + ylab("Intra-party Dimension") + theme_bw() + facet_wrap(~session, ncol = 5) + 
  geom_smooth(method = "lm", color = "black", alpha = .2, size = .5)
ggsave(filename = "plots/us/ideal_nsep_e.pdf", width = 6, height = 8)

# lm plots
gamma_est_full %>% 
  filter(parameter == "r_sq") %>% 
  select(mean, q5, q95, model, dim, session) %>% 
  ggplot(aes(x = session, y = mean, ymin = q5, ymax = q95, shape = model)) + geom_point(size = 3, position = position_dodge(width = .3)) + geom_linerange(position = position_dodge(width = .3)) + theme_bw() +
  xlab("Congress") + ylab("R^2") + facet_wrap(~dim, ncol = 1, labeller = labeller(dim = supp.labs), scales = "free") + geom_smooth(aes(linetype = model), position = position_dodge(width = .3), color = "black", size = .5, alpha = .2) + 
  theme(legend.position = "bottom", legend.title = element_blank()) + scale_linetype_discrete(labels = c("Non-separable", "Separable")) + scale_shape_discrete(labels = c("Non-separable", "Separable"))
ggsave(filename = "plots/us/rsq_e.pdf", width = 6, height = 6)

gamma_est_full %>% 
  filter(parameter != "r_sq") %>% 
  mutate(dim = rep(c(1,2), length.out = 440),
         coef = c(rep(rep(c(1,2), each = 2), length.out = 8),
                  rep(rep(c(1,2,3), each = 2), length.out = 432))) %>% 
  filter(coef > 1) %>% 
  select(mean, q5, q95, model, dim, session, coef) %>% 
  ggplot(aes(x = session, y = mean, ymin = q5, ymax = q95, shape = model)) + 
  geom_point(position = position_dodge(width = .3)) + geom_linerange(position = position_dodge(width = .3)) + theme_bw() +
  xlab("Congress") + ylab("Coefficient") + facet_grid(dim~coef, labeller = labeller(dim = supp.labs, coef = coef.labs), scales = "free") + 
  geom_smooth(aes(linetype = model), position = position_dodge(width = .3), color = "black", size = .5, alpha = .2) + 
  theme(legend.position = "bottom", legend.title = element_blank()) + scale_linetype_discrete(labels = c("Non-separable", "Separable")) + scale_shape_discrete(labels = c("Non-separable", "Separable"))
ggsave(filename = "plots/us/coef_e.pdf", width = 6, height = 6)

names(supp.labs) <- c("d1", "d2")

id_est_full %>% select(mean, model, id, dim, session) %>% 
  pivot_wider(names_from = dim, values_from = mean) %>% 
  rename(d1 = `1`, d2 = `2`) %>% 
  pivot_wider(names_from = model, values_from = c("d1", "d2")) %>% 
  group_by(session) %>% 
  summarize(d1 = summary(lm(d1_nonsep ~ d1_sep))$r.squared,
            d2 = summary(lm(d2_nonsep ~ d2_sep))$r.squared) %>% 
  ungroup() %>% 
  pivot_longer(cols = c("d1", "d2")) %>% 
  ggplot(aes(x = session, y = value)) + geom_point() + facet_wrap(~name, ncol = 1, labeller = labeller(name = supp.labs)) + geom_smooth(size = .5, color = "black", alpha = .2) + theme_bw() +
  xlab("Congress") + ylab("Correlation across Specifications")
ggsave(filename = "plots/us/cor_spec_e.pdf", width = 4, height = 4)

# loo plots
loo_est_full %>% 
  filter(elpd_diff != 0) %>% 
  ggplot(aes(x = session, y = elpd_diff, ymin = elpd_diff - 2 * se_diff, ymax = elpd_diff + 2 * se_diff, shape = model, linetype = model)) + 
  geom_linerange(position = position_dodge(width = .3)) + 
  geom_point(position = position_dodge(width = .3)) + 
  theme_bw() + xlab("Congress") + ylab("LOO Comparison") + 
  theme(legend.position = "bottom", legend.title = element_blank()) + 
  scale_linetype_discrete(labels = c("Non-separable", "Separable")) + scale_shape_discrete(labels = c("Non-separable", "Separable"))
ggsave(filename = "plots/us/loo.pdf", width = 4, height = 4)

loo_est_full %>%
  mutate(diff = elpd_diff) %>%
  filter(diff == 0) %>%
  group_by(model) %>%
  summarize(cnt = n())
## nonsep is "better" per this metric in 13 cases, sep in 24.

loo_est_full %>%
  mutate(diff = elpd_diff + 2 * se_diff) %>%
  filter(diff < 0) %>%
  group_by(model) %>%
  summarize(cnt = n())
## nonsep is "significantly better" per this metric in 1 cases, sep in 4.

# prediction and complexity
prcomp_est_full %>% 
  filter(parameter == "p_share") %>% 
  select(session, model, mean) %>% 
  pivot_wider(values_from = mean, names_from = model) %>% 
  ggplot(aes(x = nonsep, y = sep)) + 
  geom_point(alpha = .25) + 
  geom_abline(aes(slope = 1, intercept = 0)) +
  theme_bw() + xlab("Non-separable Specification") + ylab("Separable Specification") + 
  theme(legend.position = "bottom", legend.title = element_blank())
ggsave(filename = "plots/us/pred.pdf", width = 3.5, height = 3.5)

prcomp_est_full %>% 
  filter(parameter == "m_prop_compl") %>% 
  select(session, model, mean) %>% 
  pivot_wider(values_from = mean, names_from = model) %>% 
  ggplot(aes(x = nonsep, y = sep)) + 
  geom_point(alpha = .25) + 
  theme_bw() + xlab("Non-separable Specification") + ylab("Separable Specification") + 
  theme(legend.position = "bottom", legend.title = element_blank()) + 
  geom_abline(aes(slope = 1, intercept = 0))
ggsave(filename = "plots/us/compl.pdf", width = 3.5, height = 3.5)

id_est_full %>% 
  select(mean, model, id, dim, session) %>% 
  group_by(model, dim, session) %>% 
  mutate(mean = mean - mean(mean)) %>% 
  pivot_wider(id_cols = c("model", "session", "id"), names_from = "dim", values_from = "mean") %>% 
  rename(d1 = `1`, d2 = `2`) %>% 
  pivot_wider(names_from = "model", values_from = c("d1", "d2")) %>% 
  mutate(dist = ((d1_sep - d1_nonsep)^2 + (d2_sep - d2_nonsep)^2),
         dist1 = ((d1_sep - d1_nonsep)),
         dist2 = ((d2_sep - d2_nonsep))) %>% 
  merge(id_info_full, by.x = c("session", "id"), by.y = c("session", "icpsr"))  %>% 
  filter(session %in% 89:99) %>% 
  group_by(bioname) %>%
  summarize(m_d2_sep = mean(d2_sep),
            m_d2_nonsep = mean(d2_nonsep),
            m_d = (m_d2_sep - m_d2_nonsep),
            cnt = n(),
            party_code = first(party_code)) %>% 
  ungroup() %>% 
  filter(cnt > 3) %>% 
  mutate(direction = ifelse(abs(m_d2_nonsep) > abs(m_d2_sep), "more extreme", "more moderate"),
         r_sep = rank(m_d2_sep),
         r_nonsep = rank(m_d2_nonsep)) %>% 
  filter(direction == "more extreme" & party_code %in% c(100, 200)) %>% 
  top_n(10, abs(m_d)) %>% 
  mutate(party = ifelse(party_code == 200, "(R)", "(D)"),
         last = factor(paste0(bioname, " ", party)),
         last = fct_reorder(last, m_d2_nonsep)) %>% 
  ggplot() + 
  geom_segment(aes(x = last, xend = last, y = m_d2_sep, yend = m_d2_nonsep)) + theme_bw() + coord_flip() + # facet_wrap(~direction, scales = "free_y", ncol = 1) +
  geom_point(aes(x = last, y = m_d2_nonsep)) + 
  geom_point(aes(x = last, y = m_d2_sep), fill = "white", shape = 21) + 
  geom_text(aes(x = last, y = ifelse(m_d2_sep < m_d2_nonsep, m_d2_nonsep + .085 * nchar(r_nonsep), m_d2_nonsep - .085 * nchar(r_nonsep)), label = r_nonsep)) + 
  geom_text(aes(x = last, y = ifelse(m_d2_sep < m_d2_nonsep, m_d2_sep - .085 * nchar(r_sep), m_d2_sep + .085 * nchar(r_sep)), label = r_sep)) + 
  xlab("") + ylab("Intra-party Dimension")
ggsave(filename = "plots/us/move_2.pdf", width = 6, height = 3)
# 137 people

id_est_full %>% 
  select(mean, model, id, dim, session) %>% 
  group_by(model, dim, session) %>% 
  mutate(mean = mean - mean(mean)) %>% 
  pivot_wider(id_cols = c("model", "session", "id"), names_from = "dim", values_from = "mean") %>% 
  rename(d1 = `1`, d2 = `2`) %>% 
  pivot_wider(names_from = "model", values_from = c("d1", "d2")) %>% 
  mutate(dist = ((d1_sep - d1_nonsep)^2 + (d2_sep - d2_nonsep)^2),
         dist1 = ((d1_sep - d1_nonsep)),
         dist2 = ((d2_sep - d2_nonsep))) %>% 
  merge(id_info_full, by.x = c("session", "id"), by.y = c("session", "icpsr"))  %>% 
  filter(session %in% 89:99) %>% 
  group_by(bioname) %>%
  summarize(m_d2_sep = mean(d1_sep),
            m_d2_nonsep = mean(d1_nonsep),
            m_d = (m_d2_sep - m_d2_nonsep),
            cnt = n(),
            party_code = first(party_code)) %>% 
  ungroup() %>% 
  filter(cnt > 3) %>% 
  mutate(direction = ifelse(abs(m_d2_nonsep) > abs(m_d2_sep), "more extreme", "more moderate"),
         r_sep = rank(m_d2_sep),
         r_nonsep = rank(m_d2_nonsep)) %>% 
  filter(direction == "more extreme" & party_code %in% c(100, 200)) %>% 
  top_n(10, abs(m_d)) %>% 
  mutate(party = ifelse(party_code == 200, "(R)", "(D)"),
         last = factor(paste0(bioname, " ", party)),
         last = fct_reorder(last, m_d2_nonsep)) %>% 
  ggplot() + 
  geom_segment(aes(x = last, xend = last, y = m_d2_sep, yend = m_d2_nonsep)) + theme_bw() + coord_flip() + # facet_wrap(~direction, scales = "free_y", ncol = 1) +
  geom_point(aes(x = last, y = m_d2_nonsep)) + 
  geom_point(aes(x = last, y = m_d2_sep), fill = "white", shape = 21) + 
  geom_text(aes(x = last, y = ifelse(m_d2_sep < m_d2_nonsep, m_d2_nonsep + .085 * nchar(r_nonsep), m_d2_nonsep - .085 * nchar(r_nonsep)), label = r_nonsep)) + 
  geom_text(aes(x = last, y = ifelse(m_d2_sep < m_d2_nonsep, m_d2_sep - .085 * nchar(r_sep), m_d2_sep + .085 * nchar(r_sep)), label = r_sep)) + 
  xlab("") + ylab("Inter-party Dimension")
ggsave(filename = "plots/us/move_1.pdf", width = 6, height = 3)

id_est_full %>% 
  select(mean, model, id, dim, session) %>% 
  group_by(model, dim, session) %>% 
  mutate(mean = mean - mean(mean)) %>% 
  pivot_wider(id_cols = c("model", "session", "id"), names_from = "dim", values_from = "mean") %>% 
  rename(d1 = `1`, d2 = `2`) %>% 
  pivot_wider(names_from = "model", values_from = c("d1", "d2")) %>% 
  mutate(dist = sqrt((d1_sep - d1_nonsep)^2 + (d2_sep - d2_nonsep)^2)) %>% 
  group_by(session) %>% 
  summarize(dist = mean(dist, na.rm = T)) %>% 
  merge(ns_est_full %>% select(session, mean)) %>% 
  ggplot(aes(mean, dist, label = session)) + geom_point() + theme_bw() + geom_label() + xlab("Non-separability") + ylab("Mean Distance")
ggsave(filename = "plots/us/dist.pdf", width = 4, height = 4)

id_est_full %>% 
  group_by(model, dim, session) %>% 
  mutate(mean = mean - mean(mean)) %>%
  select(mean, model, id, dim, session) %>% 
  pivot_wider(id_cols = c("model", "session", "id"), names_from = "dim", values_from = "mean") %>% 
  rename(d1 = `1`, d2 = `2`) %>% 
  pivot_wider(names_from = "model", values_from = c("d1", "d2")) %>% 
  mutate(dist = sqrt((d1_sep - d1_nonsep)^2 + (d2_sep - d2_nonsep)^2)) %>% 
  ggplot(aes(d1_sep, d2_sep, xend = d1_nonsep, yend = d2_nonsep)) + geom_segment(arrow = arrow(length = unit(0.03, "npc"))) + theme_bw() + facet_wrap(~session, ncol = 5) + 
  xlab("Inter-party Dimension") + ylab("Intra-party Dimension")
ggsave(filename = "plots/us/dist2.pdf", width = 6, height = 8)